use crate::common::{config::AgentConfig, errors::Result};
use crate::common::proto::{self, oc_cluster};
use tokio_rustls::TlsConnector;
use rustls::{ClientConfig, RootCertStore, Certificate, PrivateKey};
use std::sync::Arc;
use tokio::net::TcpStream;
use tracing::{info, error, debug, warn};
use tokio::time::{sleep, Duration};

pub async fn start_tls_client(config: AgentConfig) -> Result<()> {
    let server_addr = format!("{}:{}", config.server.address, config.server.port);
    info!("Connecting to server: {}", server_addr);
    
    let mut attempt = 0;
    let max_attempts = config.agent.max_reconnect_attempts;
    let mut server_list = if let Some(list) = &config.server.addresses { list.clone() } else { vec![server_addr.clone()] };
    if !server_list.contains(&server_addr) { server_list.push(server_addr.clone()); }

    loop {
        for addr in &server_list {
            match connect_to_server(addr, &config).await {
                Ok(_) => { info!("Successfully connected to server: {}", addr); break; }
                Err(e) => { warn!("Connect to {} failed: {}", addr, e); }
            }
        }
        attempt += 1;
        if attempt >= max_attempts { return Err(anyhow::anyhow!("exceeded max attempts")); }
        sleep(Duration::from_secs(config.server.retry_interval)).await;
    }
    
    // 保持连接活跃
    loop {
        sleep(Duration::from_secs(config.agent.heartbeat_interval)).await;
        
        if let Err(e) = send_heartbeat(&server_addr, &config).await {
            warn!("Failed to send heartbeat: {}", e);
        }
    }
}

async fn connect_to_server(server_addr: &str, config: &AgentConfig) -> Result<()> {
    // 建立TCP连接
    let tcp_stream = TcpStream::connect(server_addr).await?;
    
    // 加载TLS配置
    let tls_config = load_tls_config(&config.tls)?;
    let connector = TlsConnector::from(Arc::new(tls_config));
    
    // 建立TLS连接
    let tls_stream = connector.connect(
        tokio_rustls::rustls::ServerName::try_from("localhost")?,
        tcp_stream
    ).await?;
    
    debug!("TLS connection established");
    
    // 发送注册消息（Protobuf）
    send_register_message(&mut tls_stream, config).await?;
    
    Ok(())
}

fn load_tls_config(tls_config: &crate::common::tls_config::TlsConfig) -> Result<ClientConfig> {
    let mut root_store = RootCertStore::empty();

    if let Some(ca_file) = &tls_config.ca_file {
        let ca_cert = std::fs::read(ca_file)?;
        let ca_certs = rustls_pemfile::certs(&mut &ca_cert[..])?;
        for cert in ca_certs {
            root_store.add(&Certificate(cert))?;
        }
    }

    // 如果启用客户端验证，加载客户端证书
    if tls_config.verify_client {
        let cert_bytes = std::fs::read(&tls_config.cert_file)?;
        let key_bytes = std::fs::read(&tls_config.key_file)?;
        let certs = rustls_pemfile::certs(&mut &cert_bytes[..])?
            .into_iter().map(Certificate).collect::<Vec<_>>();
        let mut keys = rustls_pemfile::pkcs8_private_keys(&mut &key_bytes[..])?;
        let key = PrivateKey(keys.remove(0));

        let config = ClientConfig::builder()
            .with_safe_defaults()
            .with_root_certificates(root_store)
            .with_client_auth_cert(certs, key)?;
        return Ok(config);
    }

    let config = ClientConfig::builder()
        .with_safe_defaults()
        .with_root_certificates(root_store)
        .with_no_client_auth();
    Ok(config)
}

async fn send_register_message(
    tls_stream: &mut tokio_rustls::client::TlsStream<TcpStream>,
    config: &AgentConfig,
) -> Result<()> {
    use prost::Message;
    let now = chrono::Utc::now().timestamp();
    let node = oc_cluster::NodeInfo {
        id: config.agent.id.clone(),
        hostname: config.agent.hostname.clone(),
        ip_address: get_local_ip().unwrap_or_else(|| "127.0.0.1".to_string()),
        tags: config.agent.tags.clone(),
        status: oc_cluster::NodeStatus::Online as i32,
        last_heartbeat: now,
        capabilities: config.agent.capabilities.clone(),
        metadata: Default::default(),
        group_ids: vec![],
    };
    let mut payload = Vec::new();
    node.encode(&mut payload)?;

    let msg = oc_cluster::ClusterMessage {
        id: uuid::Uuid::new_v4().to_string(),
        r#type: oc_cluster::MessageType::Register as i32,
        timestamp: now,
        source: config.agent.id.clone(),
        target: String::new(),
        payload,
        signature: vec![],
    };
    let mut frame = Vec::new();
    msg.encode(&mut frame)?;
    proto::write_frame(tls_stream, &frame).await?;
    info!("Registration message sent (protobuf)");
    Ok(())
}

async fn send_heartbeat(server_addr: &str, config: &AgentConfig) -> Result<()> {
    use tokio::io::AsyncWriteExt;
    use prost::Message;
    let tcp_stream = TcpStream::connect(server_addr).await?;
    let tls_config = load_tls_config(&config.tls)?;
    let connector = TlsConnector::from(Arc::new(tls_config));
    let mut tls_stream = connector.connect(
        tokio_rustls::rustls::ServerName::try_from("localhost")?,
        tcp_stream
    ).await?;

    let now = chrono::Utc::now().timestamp();
    let hb = oc_cluster::Heartbeat {
        node_id: config.agent.id.clone(),
        timestamp: now,
        health: Some(oc_cluster::NodeHealth{
            is_healthy: true,
            cpu_usage: 0.0,
            memory_usage: 0.0,
            disk_usage: 0.0,
            warnings: vec![],
            errors: vec![],
        }),
        metrics: Default::default(),
    };
    let mut hb_payload = Vec::new();
    hb.encode(&mut hb_payload)?;
    let msg = oc_cluster::ClusterMessage {
        id: uuid::Uuid::new_v4().to_string(),
        r#type: oc_cluster::MessageType::Heartbeat as i32,
        timestamp: now,
        source: config.agent.id.clone(),
        target: String::new(),
        payload: hb_payload,
        signature: vec![],
    };
    let mut frame = Vec::new();
    msg.encode(&mut frame)?;
    crate::common::proto::write_frame(&mut tls_stream, &frame).await?;
    Ok(())
}

fn get_local_ip() -> Option<String> {
    // 获取本地IP地址的简化实现
    // 实际应该使用更可靠的方法
    std::net::TcpListener::bind("0.0.0.0:0")
        .ok()
        .and_then(|listener| listener.local_addr().ok())
        .map(|addr| addr.ip().to_string())
} 